a33d35
@@ -174,9 +174,6 @@
protected void optimizeOperatorPlan(ParseContext pCtx, Set<ReadEntity> inputs,
     runStatsAnnotation(procCtx);
     perfLogger.PerfLogEnd(this.getClass().getName(), PerfLogger.TEZ_COMPILER, "Setup stats in the operator plan");
 
-    // Update bucketing version of ReduceSinkOp if needed
-    updateBucketingVersionForUpgrade(procCtx);
-
     // run Sorted dynamic partition optimization
     if(HiveConf.getBoolVar(procCtx.conf, HiveConf.ConfVars.DYNAMICPARTITIONING) &&
         HiveConf.getVar(procCtx.conf, HiveConf.ConfVars.DYNAMICPARTITIONINGMODE).equals("nonstrict") &&
@@ -229,6 +226,15 @@
protected void optimizeOperatorPlan(ParseContext pCtx, Set<ReadEntity> inputs,
       new ConstantPropagate(ConstantPropagateOption.SHORTCUT).transform(procCtx.parseContext);
     }
 
+    // ATTENTION : DO NOT, I REPEAT, DO NOT WRITE ANYTHING AFTER updateBucketingVersionForUpgrade()
+    // ANYTHING WHICH NEEDS TO BE ADDED MUST BE ADDED ABOVE
+    // This call updates the bucketing version of final ReduceSinkOp based on
+    // the bucketing version of FileSinkOp. This operation must happen at the
+    // end to ensure there is no further rewrite of plan which may end up
+    // removing/updating the ReduceSinkOp as was the case with SortedDynPartitionOptimizer
+    // Update bucketing version of ReduceSinkOp if needed
+    updateBucketingVersionForUpgrade(procCtx);
+
   }
 
   private void runCycleAnalysisForPartitionPruning(OptimizeTezProcContext procCtx,
@@ -1833,30 +1839,23 @@
private void updateBucketingVersionForUpgrade(OptimizeTezProcContext procCtx) {
 
 
     for (FileSinkOperator fsOp : fsOpsAll) {
-      Operator<?> parentOfFS = fsOp.getParentOperators().get(0);
-      if (parentOfFS instanceof GroupByOperator) {
-        GroupByOperator gbyOp = (GroupByOperator) parentOfFS;
-        List<String> aggs = gbyOp.getConf().getAggregatorStrings();
-        boolean compute_stats = false;
-        for (String agg : aggs) {
-          if (agg.equalsIgnoreCase("compute_stats")) {
-            compute_stats = true;
-            break;
-          }
-        }
-        if (compute_stats) {
+      if (!fsOp.getConf().getTableInfo().isSetBucketingVersion()) {
+        continue;
+      }
+      // Look for direct parent ReduceSinkOp
+      // If there are more than 1 parent, bail out.
+      Operator<?> parent = fsOp;
+      List<Operator<?>> parentOps = parent.getParentOperators();
+      while (parentOps != null && parentOps.size() == 1) {
+        parent = parentOps.get(0);
+        if (!(parent instanceof ReduceSinkOperator)) {
+          parentOps = parent.getParentOperators();
           continue;
         }
-      }
 
-      // Not compute_stats
-      Set<ReduceSinkOperator> rsOps = OperatorUtils.findOperatorsUpstream(parentOfFS, ReduceSinkOperator.class);
-      if (rsOps.isEmpty()) {
-        continue;
-      }
-      // Skip setting if the bucketing version is not set in FileSinkOp.
-      if (fsOp.getConf().getTableInfo().isSetBucketingVersion()) {
-        rsOps.iterator().next().setBucketingVersion(fsOp.getConf().getTableInfo().getBucketingVersion());
+        // Found the target RSOp
+        parent.setBucketingVersion(fsOp.getConf().getTableInfo().getBucketingVersion());
+        break;
       }
     }
   }
